"""Snapshot tests for RayCodeGen and SlurmCodeGen code generation.

These tests validate that RayCodeGen and SlurmCodeGen produce consistent output
for various configurations. The expected outputs are stored in:
- testdata/ray_codegen/*.py
- testdata/slurm_codegen/*.py

To update snapshots when intentional changes are made:
    UPDATE_SNAPSHOT=1 pytest tests/unit_tests/test_sky/backends/test_task_codegen.py
"""
import os
from pathlib import Path

import pytest

from sky.backends import task_codegen

TESTDATA_DIR = Path(__file__).parent / 'testdata' / 'ray_codegen'
SLURM_TESTDATA_DIR = Path(__file__).parent / 'testdata' / 'slurm_codegen'


def assert_codegen_matches_snapshot(test_name: str,
                                    generated_code: str,
                                    testdata_dir: Path = TESTDATA_DIR) -> None:
    """Compare generated code against snapshot file.

    Args:
        test_name: Name of the test (used to find snapshot file)
        generated_code: The code generated by RayCodeGen or SlurmCodeGen
        testdata_dir: Directory containing test snapshots (default: ray_codegen)

    If UPDATE_SNAPSHOT=1 env var is set, updates the snapshot file instead
    of comparing.
    """
    snapshot_path = testdata_dir / f'{test_name}.py'

    if os.environ.get('UPDATE_SNAPSHOT') == '1':
        # Update snapshot
        snapshot_path.parent.mkdir(parents=True, exist_ok=True)
        snapshot_path.write_text(generated_code)
        print(f'Updated snapshot: {snapshot_path}')
        return

    # Read expected snapshot
    if not snapshot_path.exists():
        pytest.fail(f'Snapshot file not found: {snapshot_path}\n'
                    f'Run with UPDATE_SNAPSHOT=1 to create it.')

    expected = snapshot_path.read_text()

    # Compare
    if generated_code != expected:
        # Show helpful diff message
        import difflib
        diff = difflib.unified_diff(
            expected.splitlines(keepends=True),
            generated_code.splitlines(keepends=True),
            fromfile=f'{test_name}.py (expected)',
            tofile=f'{test_name}.py (actual)',
        )
        diff_text = ''.join(diff)
        pytest.fail(
            f'Generated code does not match snapshot: {snapshot_path}\n\n'
            f'Diff:\n{diff_text}\n\n'
            f'Run with UPDATE_SNAPSHOT=1 to update the snapshot.')


def test_single_node_with_gpu():
    """Test single-node GPU job with setup commands.

    Mirrors usage in _execute_task_one_node() with GPU and setup.
    """
    codegen = task_codegen.RayCodeGen()
    codegen.add_prologue(job_id=2)

    resources_dict = {'CPU': 4.0, 'GPU': 1.0}
    task_env_vars = {
        'SKYPILOT_TASK_ID': 'sky-2024-11-17-00-00-00-000001-cluster-2',
        'MODEL_NAME': 'resnet50',
    }

    codegen.add_setup(
        1,
        resources_dict=resources_dict,
        stable_cluster_internal_ips=['10.0.0.1'],
        env_vars=task_env_vars,
        setup_cmd='pip install torch',
        setup_log_path='/sky/logs/setup.log',
    )

    # Single node: no gang_scheduling_id passed
    codegen.add_task(
        1,
        bash_script='python train.py',
        task_name='train_task',
        resources_dict={
            'CPU': 4.0,
            'GPU': 1.0
        },
        log_dir='/sky/logs/tasks',
        env_vars=task_env_vars,
    )

    codegen.add_epilogue()

    result = codegen.build()
    assert_codegen_matches_snapshot('single_node_with_gpu', result)


def test_multi_node_2nodes():
    """Test multi-node job (2 nodes).

    Mirrors usage in _execute_task_n_nodes() for basic multi-node setup.
    """
    codegen = task_codegen.RayCodeGen()
    codegen.add_prologue(job_id=3)

    num_nodes = 2
    resources_dict = {'CPU': 2.0}
    task_env_vars = {
        'SKYPILOT_TASK_ID': 'sky-2024-11-17-00-00-00-000002-cluster-3',
    }

    codegen.add_setup(
        num_nodes,
        resources_dict=resources_dict,
        stable_cluster_internal_ips=['10.0.0.1', '10.0.0.2'],
        env_vars=task_env_vars,
        setup_cmd=None,
        setup_log_path='/sky/logs/setup.log',
    )

    codegen.add_task(
        num_nodes,
        bash_script='echo "Running on node $SKYPILOT_NODE_RANK"',
        task_name='distributed_task',
        resources_dict={'CPU': 2.0},
        log_dir='/sky/logs/tasks',
        env_vars=task_env_vars,
    )

    codegen.add_epilogue()

    result = codegen.build()
    assert_codegen_matches_snapshot('multi_node_2nodes', result)


def test_slurm_single_node_with_gpu():
    """Test single-node GPU job with setup commands on Slurm.

    Mirrors test_single_node_with_gpu() but for SlurmCodeGen.
    """
    codegen = task_codegen.SlurmCodeGen(slurm_job_id='12345')
    codegen.add_prologue(job_id=2)

    resources_dict = {'CPU': 4.0, 'GPU': 1.0}
    task_env_vars = {
        'SKYPILOT_TASK_ID': 'sky-2024-11-17-00-00-00-000001-cluster-2',
        'MODEL_NAME': 'resnet50',
    }

    codegen.add_setup(
        1,
        resources_dict=resources_dict,
        stable_cluster_internal_ips=['10.0.0.1'],
        env_vars=task_env_vars,
        setup_cmd='pip install torch',
        setup_log_path='/sky/logs/setup.log',
    )

    codegen.add_task(
        1,
        bash_script='python train.py',
        task_name='train_task',
        resources_dict={
            'CPU': 4.0,
            'GPU': 1.0
        },
        log_dir='/sky/logs/tasks',
        env_vars=task_env_vars,
    )

    codegen.add_epilogue()

    result = codegen.build()
    assert_codegen_matches_snapshot('slurm_single_node_with_gpu',
                                    result,
                                    testdata_dir=SLURM_TESTDATA_DIR)
